DEBUG = False

class MarkovChain(object):
    def __init__(self, states, start_probabilities, state_transitions):
        observations_prob = {}
        for state in states:
            observations_prob[state] = {state:1}
            
        self.hmm = MyHMM(states, states, start_probabilities, state_transitions, observations_prob)

    def train(self,state_sequences, normalize = False):
        self.hmm.train(state_sequences, [], normalize)
        observations_prob = {}
        for state in self.hmm.states:
            observations_prob[state] = {state:1}
        self.hmm.observations_prob=observations_prob
        self.hmm.observations = self.hmm.states
    
    def probability(self, sequence):
        return self.hmm.forward(sequence)
    

            
class MyHMM(object):
    def __init__(self, states, observations, start_probabilities, state_transitions, observations_prob, coef=4.74081266062e-008):
        """coef is used to allow for unseen data. Use the probability of a word that has one or two occurences"""
        self.states = states
        self.observations = observations
        self.start_probabilities = start_probabilities
        self.state_transitions = state_transitions
        self.observations_prob = observations_prob
        self.coef=coef
        self.len_states = len(self.states)
        self.len_observations = len(self.observations)

    def __str__(self):
        return 'MyHMM with %d states and %d observations' %(self.len_states, len(self.observations))
    
    def train(self, state_sequences, state_observations, normalize = False): 
        self.train_states(state_sequences, normalize)
        self.train_observations(state_observations, normalize)
        
    def update_self(self):
        self.len_states = len(self.states)
        self.len_observations = len(self.observations)
                
        

    def train_states(self, state_sequences, normalize):
        state_count = 0      
        start_count = 0
        for sequence,prob in state_sequences:
            prev = None
            for part in sequence:
                if not part in self.states: self.states.append(part)
                if prev:
                    self.state_transitions[prev] = self.state_transitions.get(prev, {})
                    self.state_transitions[prev][part] = self.state_transitions[prev].get(part, 0) + prob  
                    state_count += prob
                    if not normalize:
                        assert(self.state_transitions[prev][part]<=1)              
                else:
                    self.start_probabilities[part] = self.start_probabilities.get(part, 0) + prob
                    start_count += prob
                    if not normalize:
                        assert(self.start_probabilities[part]<=1)
                prev = part
                
        if normalize:
            for key in self.state_transitions:
                for key2 in self.state_transitions[key]:
                    self.state_transitions[key][key2] = float(self.state_transitions[key][key2])/state_count
                    assert(self.state_transitions[key][key2]<=1) 
            for key in self.start_probabilities:
                self.start_probabilities[key]= float(self.start_probabilities[key])/start_count
                assert(self.start_probabilities[key]<=1) 
        self.update_self()

    def train_observations(self, state_observations, normalize, addStates = False):
        obs_count = 0
        for sequence,prob in state_observations:
            for (observation, state) in sequence:
                if not state in self.states:
                    if not addStates: print state.encode('utf-8'), 'NOT FOUND IN STATES!!!!!!!!'
                    else: self.states.append(state)
                if not observation in self.observations: self.observations.append(observation)
                self.observations_prob[state] = self.observations_prob.get(state, {})
                self.observations_prob[state][observation] = self.observations_prob[state].get(observation,0) + prob
                obs_count += prob
                if not normalize:
                        assert(self.observations_prob[state][observation]<=1)
        if normalize:
                for key in self.observations_prob:
                    for key2 in self.observations_prob[key]:
                        self.observations_prob[key][key2] = float(self.observations_prob[key][key2])/obs_count
                        assert(self.observations_prob[key][key2]<=1) 
        self.update_self()                        
        
    def viterbi(self, sequence):
        global DEBUG
        d = []
        f = []
        coef = self.coef
        for i in range(0, len(sequence)):
            d.append([])
            f.append([])
        sequence_0=sequence[0]
        for state_i in self.states:
#            state_i=self.states[i]
            try:
                a = self.start_probabilities[state_i]
            except KeyError:
                a = coef
                self.start_probabilities[state_i] = coef
                if DEBUG: print 'Start probability not found:',state_i.encode('utf-8')
            try:
                b = self.observations_prob[state_i]
            except KeyError:
                b= {}
                self.observations_prob[state_i] = b
                if DEBUG: 'Observation for state not found:', state_i.encode('utf-8')
            try:
                c = b[sequence_0]
            except KeyError:
                c = coef
                b[sequence_0] = coef
                if DEBUG: 'Observation probability for state, part not found:',state_i.encode('utf-8'),sequence_0
                
            d[0].append(a*c)
        if len(sequence)>1:
            for t in xrange(1, len(sequence)): # for each t > 1
                state_i = None
                sequence_t = sequence[t]
                for state_i in self.states:
                    dtemp = []
                    ftemp = []
                    for j,state_j in enumerate(self.states):
                        try:
                            a = self.state_transitions[state_j]
                        except KeyError:
                            a={}
                            self.state_transitions[state_j]=a
                            if DEBUG: print 'State transition dict for state not found:',state_j.encode('utf-8')
                        try:
                            b = a[state_i]
                        except KeyError:
                            b = coef
                            a[state_i]=coef
                            if DEBUG: print 'State transition for state, state not found:',state_j.encode('utf-8'),state_i.encode('utf-8')
                        
                        daji = d[t-1][j] * b
                        
                        try:
                            c = self.observations_prob[state_i]
                        except KeyError:
                            c = {}
                            self.observations_prob[state_i] = c
                            if DEBUG: assert self.observations_prob[state_i] == {}
                            if DEBUG: 'Observation for state not found:', state_i.encode('utf-8')                       
                        try:
                            e = c[sequence_t]
                        except KeyError:
                            e = coef
                            c[sequence_t]=coef
                            if DEBUG: 'Observation probability for state, part not found:',state_i.encode('utf-8'),sequence_t                       
                            
                        dtemp.append(daji * e)
                        ftemp.append(daji)
        
                    f[t].append(ftemp.index(max(ftemp)))
                    d[t].append(max(dtemp))
        else:
            t=0
        state_seq = []
        it = d[t].index(max(d[t]))
        state_seq.append(self.states[it])
        for t in reversed(range(0, len(sequence)-1)):
            it = f[t+1][it]
            state_seq.append(self.states[it])
                          
    #    for t in range(0,len(sequence)):
    #        state_seq.append(states[f[t]])
        state_seq.reverse()
        return state_seq
                
    
          
    
    def forward(self, sequence):
        obs_prob_get = self.observations_prob.get

        a = []
        for i in range(0, len(sequence)): # initialize a
            a.append([])
        sequence_0 = sequence[0]
        for state_i in self.states: # for t = 1
            try:
                start_prob_i = self.start_probabilities[state_i]
            except KeyError:
                start_prob_i = self.coef
                self.start_probabilities[state_i] = start_prob_i
            try:
                obs_prob_i = self.observations_prob[state_i]
            except KeyError:
                obs_prob_i = {}
                self.observations_prob[state_i] = obs_prob_i
            try:
                obs_prob_i_0 = obs_prob_i[sequence_0]
            except KeyError:
                obs_prob_i_0 = self.coef
                obs_prob_i[sequence_0] = obs_prob_i_0
            a[0].append(start_prob_i * obs_prob_i_0)
        for t in xrange(1, len(sequence)): # for each t > 1
            sequence_t=sequence[t]
            for state_j in self.states: # for each state
                suma = 0
                for i, state_i in enumerate(self.states):
                    ati = a[t-1][i]
                    try:
                        state_trans_i = self.state_transitions[state_i]
                    except KeyError:
                        state_trans_i = {}
                        self.state_transitions[state_i] = state_trans_i
                    try:
                        state_trans_i_j = state_trans_i[state_j]
                    except KeyError:
                        state_trans_i_j = self.coef    
                        state_trans_i[state_j] = state_trans_i_j
                    aij = state_trans_i_j
                    suma=suma + ati*aij# SUM[at(i)*aij]
                try:
                    obs_prob_j = self.observations_prob[state_j]
                except KeyError:
                    obs_prob_j = {}
                    self.observations_prob[state_j] = obs_prob_j
                try:
                    obs_prob_j_t = obs_prob_j[sequence_t]
                except KeyError:
                    obs_prob_j_t = self.coef
                    obs_prob_j[sequence_t] = obs_prob_j_t
                bjkt = obs_prob_j_t
                a[t].append(bjkt * suma)
        total_prob = 0
        t = len(sequence)-1
        for j in xrange(0, self.len_states):
            total_prob+= a[t][j]
        return total_prob
    
if __name__ == '__main__':
    from decimal import Decimal
    obs_sequence = ['dry', 'damp', 'soggy', 'dry']
    state_transitions = {
     'sun':{'sun':Decimal('0.5'), 'cloud':Decimal('0.25'), 'rain':Decimal('0.25')}, 
     'cloud':{'sun':Decimal('0.375'), 'cloud':Decimal('0.125'), 'rain':Decimal('0.375')}, 
     'rain':{'sun':Decimal('0.125'), 'cloud':Decimal('0.675'), 'rain':Decimal('0.375')}
                     }     
    start_probabilities = {'sun':Decimal('0.63'), 'cloud':Decimal('0.17'), 'rain':Decimal('0.20')}
    
    states = ['sun', 'cloud', 'rain']
    observations = ['dry', 'dryish', 'damp', 'soggy']
    
    observations_prob = {
    'sun':{'dry':Decimal('0.60'), 'dryish':Decimal('0.20'), 'damp':Decimal('0.15'), 'soggy':Decimal('0.05')}, 
    'cloud':{'dry':Decimal('0.25'), 'dryish':Decimal('0.25'), 'damp':Decimal('0.25'), 'soggy':Decimal('0.25')}, 
    'rain':{'dry':Decimal('0.05'), 'dryish':Decimal('0.10'), 'damp':Decimal('0.35'), 'soggy':Decimal('0.50')}
    }
    
    hmm = MyHMM(states, observations, start_probabilities, state_transitions, observations_prob)
    
    #print getcontext()

    print hmm.forward(obs_sequence)
    assert str(hmm.forward(obs_sequence))=='0.0074792177734375000'
    print 'viterbi:'
    print hmm.viterbi(obs_sequence)
    assert hmm.viterbi(obs_sequence)==['sun', 'rain', 'cloud', 'sun']
        